// @Version: 1.0.0
// @Author: Bruce
// @TIME  : 2020/3/2 20:26
// @E-mail: mzpy_1119@126.com

package main

import (
	"fmt"
	"github.com/BurntSushi/toml"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"net"
	"strconv"
)

type tomlConf struct {
	//Title    string `toml:"title"`
	SSH      sshConf
	SSH_Bind sshBindConf
}

type sshConf struct {
	Host string `toml:"ssh_host"`
	Port int    `toml:"ssh_port"`
	User string `toml:"ssh_user"`
	Pkey string `toml:"ssh_pkey"`
	Pwd  string `toml:"ssh_pass"`
}

type sshBindConf struct {
	RemoteAddr   string     `toml:"remote_bind_addr"`
	LocalAddr    string     `toml:"local_bind_addr"`
	BindPortList []BindPort `toml:"bind_port"`
}

type BindPort struct {
	Remote string
	Port   string
}

type confs struct {
	List []tomlConf `toml:"item"`
}

func main() {
	var filePath = "./conf/app.toml"
	//var config tomlConf
	var configs confs
	if _, err := toml.DecodeFile(filePath, &configs); err != nil {
		fmt.Println(err)
		return
	}
	//fmt.Println(configs)
	for _, config := range configs.List {
		go conn(config)
	}
	select {}
}

func conn(config tomlConf) {
	conf := &ssh.ClientConfig{
		User: config.SSH.User,
		Auth: []ssh.AuthMethod{
			//ssh.PublicKeys(signer),
			ssh.Password(config.SSH.Pwd),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	check(conf, config)

	// Loop local port array
	for _, p := range config.SSH_Bind.BindPortList {
		go func(port BindPort) {
			localAddr := fmt.Sprintf("%s:%s", config.SSH_Bind.LocalAddr, port.Port)
			fmt.Print("Listen success: ", localAddr, "\n")

			localListener, err := net.Listen("tcp", localAddr)

			if err != nil {
				log.Fatalf("net.Listen failed: %v", err)
			}

			for {
				// Setup localConn (type net.Conn)
				localConn, err := localListener.Accept()
				if err != nil {
					log.Fatalf("listen.Accept failed: %v", err)
				}
				go forward(localConn, conf, config, port.Remote)
			}
		}(p)
	}

}

func check(cliconfig *ssh.ClientConfig, config tomlConf) {
	//dial 获取ssh client
	addr := fmt.Sprintf("%s:%d", config.SSH.Host, config.SSH.Port)
	sshClient, err := ssh.Dial("tcp", addr, cliconfig)
	if err != nil {
		log.Fatal("创建ssh client 失败", err)
	}
	defer sshClient.Close()

	//创建ssh-session
	session, err := sshClient.NewSession()
	if err != nil {
		log.Fatal("创建ssh session 失败", err)
	}
	defer session.Close()
	//执行远程命令
	combo, err := session.CombinedOutput("whoami;uptime;df -h; free -m;top -bn 3 | grep 'Cpu'")
	if err != nil {
		log.Fatal("远程执行cmd 失败", err)
	}
	log.Println("命令输出:", string(combo))

}

func forward(localConn net.Conn, cliconfig *ssh.ClientConfig, config tomlConf, port string) {

	sshAddr := fmt.Sprintf("%s:%s", config.SSH.Host, strconv.Itoa(config.SSH.Port))
	//bindAddr := fmt.Sprintf("%s:%s", config.SSH_Bind.RemoteAddr, strconv.Itoa(port))

	// Setup sshClientConn (type *ssh.ClientConn)
	sshClientConn, err := ssh.Dial("tcp", sshAddr, cliconfig)
	if err != nil {
		log.Fatalf("ssh.Dial failed: %s", err)
	}

	// Setup sshConn (type net.Conn)
	//sshConn, err := sshClientConn.Dial("tcp", bindAddr)
	sshConn, err := sshClientConn.Dial("tcp", port)
	// Copy localConn.Reader to sshConn.Writer
	go func() {
		_, err = io.Copy(sshConn, localConn)
		if err != nil {
			log.Fatalf("io.Copy failed: %v", err)
		}
	}()

	// Copy sshConn.Reader to localConn.Writer
	go func() {
		_, err = io.Copy(localConn, sshConn)
		if err != nil {
			log.Fatalf("io.Copy failed: %v", err)
		}
	}()
}
